package com.jd.genie.agent.tool.common;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.jd.genie.agent.agent.AgentContext;
import com.jd.genie.agent.dto.DeepSearchRequest;
import com.jd.genie.agent.dto.DeepSearchrResponse;
import com.jd.genie.agent.dto.FileRequest;
import com.jd.genie.agent.tool.BaseTool;
import com.jd.genie.agent.util.SpringContextHolder;
import com.jd.genie.agent.util.StringUtil;
import com.jd.genie.config.GenieConfig;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.context.ApplicationContext;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

@Slf4j
@Data
public class ChatBiQueryTool implements BaseTool {

    private AgentContext agentContext;

    @Override
    public String getName() {
        return "chatbi_query";
    }

    @Override
    public String getDescription() {
        String desc = "这是一个使用BI查询数据的工具，可以根据用户的自然语言查询生成SQL语句，执行查询并返回查询结果";
        GenieConfig genieConfig = SpringContextHolder.getApplicationContext().getBean(GenieConfig.class);
        return genieConfig.getChatBiQueryToolDesc().isEmpty() ? desc : genieConfig.getChatBiQueryToolDesc();
    }

    @Override
    public Map<String, Object> toParams() {

        GenieConfig genieConfig = SpringContextHolder.getApplicationContext().getBean(GenieConfig.class);
        if (!genieConfig.getChatBiQueryToolPamras().isEmpty()) {
            return genieConfig.getChatBiQueryToolPamras();
        }

        Map<String, Object> taskParam = new HashMap<>();
        taskParam.put("type", "object");
        taskParam.put("description", "需要查询的query");
        Map<String, Object> properties = new HashMap<>();
        Map<String, Object> queryParam = new HashMap<>();
        queryParam.put("type", "object");
        properties.put("query", taskParam);
        queryParam.put("properties", properties);
        queryParam.put("required", Collections.singletonList("query"));
        return queryParam;
    }

    @Override
    public Object execute(Object input) {

        try {
            ApplicationContext applicationContext = SpringContextHolder.getApplicationContext();
            GenieConfig genieConfig = applicationContext.getBean(GenieConfig.class);
            
            Map<String, Object> inputMap = (Map<String, Object>) input;
            String query = (String) inputMap.get("query");
            
            DeepSearchRequest deepSearchRequest = DeepSearchRequest.builder()
                    .request_id(agentContext.getRequestId()+ "_" + StringUtil.generateRandomString(5))
                    .query(query)
                    .agent_id("1")
                    .optional_configs(new HashMap<>())
                    .src_configs(new HashMap<>())
                    .scene_type("auto_agent")
                    .stream(true)
                    .content_stream(true)
                    .build();
            
            Future future = callChatBiQueryStream(deepSearchRequest);
            return future.get(genieConfig.getChatBiQueryToolTimeoutSec(), TimeUnit.SECONDS);
        } catch (Exception e) {
            log.error("{} chatbi_query execute error", agentContext.getRequestId(), e);
            return "BI查询执行失败: " + e.getMessage();
        }
    }

    public CompletableFuture<String> callChatBiQueryStream(DeepSearchRequest deepSearchRequest) {
        CompletableFuture<String> future = new CompletableFuture<>();
        
        try {
            ApplicationContext applicationContext = SpringContextHolder.getApplicationContext();
            GenieConfig genieConfig = applicationContext.getBean(GenieConfig.class);
            String url = genieConfig.getChatBiQueryUrl() + "/v1/tool/chatbi_query";
            OkHttpClient client = new OkHttpClient.Builder()
                    .connectTimeout(genieConfig.getChatBiQueryToolConnectTimeoutSec(), TimeUnit.SECONDS)
                    .readTimeout(genieConfig.getChatBiQueryToolReadTimeoutSec(), TimeUnit.SECONDS)
                    .writeTimeout(genieConfig.getChatBiQueryToolWriteTimeoutSec(), TimeUnit.SECONDS)
                    .callTimeout(genieConfig.getChatBiQueryToolTimeoutSec(), TimeUnit.SECONDS)
                    .build();
            
            String requestBody = JSON.toJSONString(deepSearchRequest);
            RequestBody body = RequestBody.create(
                    MediaType.parse("application/json"),
                    requestBody
            );

            log.info("===== {} =====\n === chatbi_query request: ==== \n {} \n ================ \n", agentContext.getRequestId(), requestBody);
            
            Request.Builder requestBuilder = new Request.Builder()
                    .url(url)
                    .post(body);
            Request request = requestBuilder.build();
            
            int firstInterval = genieConfig.getChatBiQueryToolFirstInterval();
            int sendInterval = genieConfig.getChatBiQueryToolSendInterval();
            
            client.newCall(request).enqueue(new Callback() {
                @Override
                public void onFailure(Call call, IOException e) {
                    log.error("{} chatbi_query request failed", agentContext.getRequestId(), e);
                    future.completeExceptionally(e);
                }
                
                @Override
                public void onResponse(Call call, Response response) throws IOException {
                    
                    log.info("===== {} =====\n ==== chatbi query response: ==== \n {} {} {} \n ================ \n", agentContext.getRequestId(), response, response.code(), response.body());
                    try (ResponseBody responseBody = response.body()) {
                        if (!response.isSuccessful() || response.body() == null) {
                            log.error("{} chatbi_query request failed with code: {}", agentContext.getRequestId(), response.code());
                            future.completeExceptionally(new RuntimeException("HTTP " + response.code()));
                            return;
                        }
                        
                        int index = 1;
                        StringBuilder stringBuilderIncr = new StringBuilder();
                        StringBuilder stringBuilderAll = new StringBuilder();
                        String line;
                        String digitalEmployee = agentContext.getToolCollection().getDigitalEmployee(getName());
                        String result = "BI查询结果为空";
                        String messageId = "";
                        
                        try (BufferedReader reader = new BufferedReader(new InputStreamReader(responseBody.byteStream(), "UTF-8"))) {
                        
                        while ((line = reader.readLine()) != null) {
                            if (line.startsWith("data: ")) {
                                String data = line.substring(6);
                                if (data.equals("[DONE]")) {
                                    break;
                                }
                                if (data.startsWith("heartbeat")) {
                                    continue;
                                }
                                
                                if (index == 1 || index % 100 == 0) {
                                    log.info("==== {} =====\n ==== chatbi_query recv data: ==== \n {} \n ================ \n", agentContext.getRequestId(), data);
                                }
                                
                                DeepSearchrResponse queryResponse = JSONObject.parseObject(data, DeepSearchrResponse.class);
                                if (queryResponse == null) {
                                    log.warn("{} Failed to parse DeepSearchrResponse from data: {}", agentContext.getRequestId(), data);
                                    continue;
                                }
                                FileTool fileTool = new FileTool();
                                fileTool.setAgentContext(agentContext);
                                
                                if (queryResponse.getIsFinal()) {
                                    if (agentContext.getIsStream()) {
                                        queryResponse.setAnswer(stringBuilderAll.toString());
                                    }
                                    if (queryResponse.getAnswer() == null || queryResponse.getAnswer().isEmpty()) {
                                        log.error("{} chatbi_query answer empty or null", agentContext.getRequestId());
                                        break;
                                    }
                                    
                                    String fileName = StringUtil.removeSpecialChars(queryResponse.getQuery() + "的搜索结果.md");
                                    String answer = queryResponse.getAnswer();
                                    String fileDesc = (answer != null && answer.length() > 0) ? 
                                            answer.substring(0, Math.min(answer.length(), genieConfig.getChatBiQueryToolFileDescTruncateLen())) + "..." : 
                                            "BI查询结果";
                                    FileRequest fileRequest = FileRequest.builder()
                                            .requestId(agentContext.getRequestId())
                                            .fileName(fileName)
                                            .description(fileDesc)
                                            .content(queryResponse.getAnswer())
                                            .build();
                                    fileTool.uploadFile(fileRequest, false, false);
                                    
                                    // 注意: DeepSearchrResponse不包含ChartInfo和QuerySql字段
                                    // 如果需要这些功能，请使用原始的ChatBiQueryResponse或扩展DeepSearchrResponse
                                    
                                    result = (answer != null && answer.length() > 0) ? 
                                            answer.substring(0, Math.min(answer.length(), genieConfig.getChatBiQueryToolMessageTruncateLen())) : 
                                            "BI查询结果为空";
                                    
                                    agentContext.getPrinter().send(messageId, "chatbi_query", queryResponse, digitalEmployee, true);
                                    
                                } else {
                                    Map<String, Object> contentMap = new HashMap<>();
                                    for (int idx = 0; idx < queryResponse.getSearchResult().getQuery().size(); idx++) {
                                        contentMap.put(queryResponse.getSearchResult().getQuery().get(idx), queryResponse.getSearchResult().getDocs().get(idx));
                                    }
                                    if ("extend".equals(queryResponse.getMessageType())) {
                                        messageId = StringUtil.getUUID();
                                        queryResponse.setSearchFinish(false);
                                        // 注意: DeepSearchrResponse没有setQueryFinish方法
                                        agentContext.getPrinter().send(messageId, "chatbi_query", queryResponse, digitalEmployee, true);
                                    } else if ("search".equals(queryResponse.getMessageType())) {
                                        queryResponse.setSearchFinish(true);
                                        agentContext.getPrinter().send(messageId, "chatbi_query", queryResponse, digitalEmployee, true);
                                        FileRequest fileRequest = FileRequest.builder()
                                                .requestId(agentContext.getRequestId())
                                                .fileName(queryResponse.getQuery() + "_search_result.txt")
                                                .description(queryResponse.getQuery() + "...")
                                                .content(JSON.toJSONString(contentMap))
                                                .build();
                                        fileTool.uploadFile(fileRequest, false, true);
                                    } else if ("report".equals(queryResponse.getMessageType())) {
                                        if (index == 1) {
                                            messageId = StringUtil.getUUID();
                                        }
                                        stringBuilderIncr.append(queryResponse.getAnswer());
                                        stringBuilderAll.append(queryResponse.getAnswer());
                                        if (index == firstInterval || index % sendInterval == 0) {
                                            queryResponse.setAnswer(stringBuilderIncr.toString());
                                            agentContext.getPrinter().send(messageId, "chatbi_query", queryResponse, digitalEmployee, false);
                                            stringBuilderIncr.setLength(0);
                                        }
                                        index++;
                                    }
                                }
                            }
                        }
                        future.complete(result);
                        
                        } catch (IOException e) {
                            log.error("{} chatbi_query reader error", agentContext.getRequestId(), e);
                            future.completeExceptionally(e);
                        }
                        
                    } catch (Exception e) {
                        log.error("{} chatbi_query request error", agentContext.getRequestId(), e);
                        future.completeExceptionally(e);
                    }
                }
            });
        } catch (Exception e) {
            log.error("{} chatbi_query request error", agentContext.getRequestId(), e);
            future.completeExceptionally(e);
        }
        
        return future;
    }
}